import os
# pip install imageAI
from imageai.Detection import ObjectDetection


def object_detect_1():
    execution_path = os.getcwd()
    detector = ObjectDetection()
    detector.setModelTypeAsRetinaNet()
    # https://github.com/OlafenwaMoses/ImageAI/releases/download/1.0/resnet50_coco_best_v2.0.1.h5
    detector.setModelPath(os.path.join(execution_path, "resnet50_coco_best_v2.0.1.h5"))
    detector.loadModel()
    detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path, "image.jpg"),
                                                 output_image_path=os.path.join(execution_path, "imagenew.jpg"))
    for eachObject in detections:
        print(eachObject["name"] + " : " + eachObject["percentage_probability"])


# https://zhuanlan.zhihu.com/p/150880766
def object_detect_2():
    execution_path = os.getcwd()
    detector = ObjectDetection()
    detector.setModelTypeAsTinyYOLOv3()
    # https://link.zhihu.com/?target=https%3A//github.com/OlafenwaMoses/ImageAI/releases/download/1.0/yolo-tiny.h5
    detector.setModelPath(os.path.join(execution_path, "yolo-tiny.h5"))
    detector.loadModel()
    detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path, "image.jpg"),
                                                 output_image_path=os.path.join(execution_path, "imagenew.jpg"))
    for eachObject in detections:
        print(eachObject["name"] + " : " + eachObject["percentage_probability"])


#     # https://blog.csdn.net/weixin_42652125/article/details/82312739
def object_detect_3():
    execution_path = os.getcwd()
    detector = ObjectDetection()
    detector.setModelTypeAsYOLOv3()
    detector.setModelPath(os.path.join(execution_path, "Fyolo.h5"))
    detector.loadModel()
    detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path, "image.jpg"),
                                                 output_image_path=os.path.join(execution_path, "imagenew.jpg"))
    for eachObject in detections:
        print(eachObject["name"] + " : " + eachObject["percentage_probability"])


if __name__ == '__main__':
    object_detect_1()
    object_detect_2()
    object_detect_3()
